from typing import Tuple, Optional, Any, Type

import numpy as np
import torch
from torch import nn

from algorithms.common import LayerNorm2d


class PromptEncoder(nn.Module):
    def __init__(
            self,
            embed_dim: int,
            image_embedding_size: Tuple[int, int],
            input_image_size: Tuple[int, int],
            mask_in_chans: int,
            activation: Type[nn.Module] = nn.GELU,
    ) -> None:
        super().__init__()
        self.embed_dim = embed_dim
        self.input_image_size = input_image_size
        self.image_embedding_size = image_embedding_size  # 来自image_size//vit_patch_size
        self.pe_layer = PositionEmbeddingRandom(embed_dim // 2)  # 这里除以2是因为有x和y两个方向，它们最终会合回来
        # 这里的值为 4
        # 第0个point_embedding的weight，给背景
        # 第1个point_embedding的weight给目标点，即点在目标范围内的点
        # 第2个point_embedding的weight给box左上角的point，第3个给右下角.
        self.num_point_embedding: int = 4
        point_embeddings = [nn.Embedding(1, embed_dim) for i in range(self.num_point_embedding)]
        self.point_embeddings = nn.ModuleList(point_embeddings)
        self.not_a_point_embed = nn.Embedding(1, embed_dim)
        self.mask_input_size = (4 * image_embedding_size[0], 4 * image_embedding_size[1])
        self.mask_downscaling = nn.Sequential(
            nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2),
            LayerNorm2d(mask_in_chans // 4),
            activation(),
            nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2),
            LayerNorm2d(mask_in_chans),
            activation(),
            nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1)
        )
        self.no_mask_embed = nn.Embedding(1, embed_dim)

    def forward(
            self,
            points: Optional[Tuple[torch.Tensor, torch.Tensor]],
            boxes: Optional[torch.Tensor],
            masks: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        不同类型的prompts，返回sparse embedding和dense embedding（稀疏嵌入和密集嵌入）
        Args:
          points Optional[Tuple[torch.Tensor, torch.Tensor]]，经过embedding的point坐标，包含label，
          boxes (torch.Tensor or none):经过embedding的boxes
          masks (torch.Tensor or none):经过embedding的masks
        :return:
          torch.Tensor:Sparse embedding for the points and boxes,with shape BxNx(embed_dim),
            where N is determined by the number of input points and boxes.
          torch.Tensor:dense embeddings for the masks,in the shape Bx(embed_dim)x(embed_H)x(embed_W)
        """
        bs = self._get_batch_size(points, boxes, masks)
        # 创建一个Tensor，张量中的数值是随机“垃圾”数据，没有经过初始化
        sparse_embeddings = torch.empty((bs, 0, self.embed_dim), device=self._get_device())
        if points is not None:
            coords, labels = points
            point_embeddings = self._embed_points(coords, labels, pad=(boxes is None))
            sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
        if boxes is not None:
            box_embeddings = self._embed_boxes(boxes)
            sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
        if masks is not None:
            dense_embeddings = self._embed_masks(masks)  # 这里的masks一般来自上一次预测，否则使用else代码段
        else:
            # no_mask_embed的维度是(1,embed_dim)
            dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(
                bs, -1, self.image_embedding_size[0], self.image_embedding_size[1]
            )
        return sparse_embeddings, dense_embeddings

    def _get_batch_size(
            self,
            points: Optional[Tuple[torch.Tensor, torch.Tensor]],
            boxes: Optional[torch.Tensor],
            masks: Optional[torch.Tensor],
    ) -> int:
        """
        Gets the batch size of the output given the batch size of the input prompts
        """
        if points is not None:
            return points[0].shape[0]
        elif boxes is not None:
            return boxes.shape[0]
        elif masks is not None:
            return masks.shape[0]
        else:
            return 1

    def _get_device(self) -> torch.device:
        return self.point_embeddings[0].weight.device

    def get_dense_pe(self) -> torch.Tensor:
        # image_embedding_size - (image_size//vit_patch_size,image_size//vit_patch_size)
        return self.pe_layer(self.image_embedding_size).unsqueeze(0)

    def _embed_points(
            self,
            points: torch.Tensor,
            labels: torch.Tensor,
            pad: bool,
    ) -> torch.Tensor:
        """Embeds point prompts."""
        points = points + 0.5  # Shift to center of pixel
        if pad:
            padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)
            padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)
            points = torch.cat([points, padding_point], dim=1)
            labels = torch.cat([labels, padding_label], dim=1)

        # point_embedding — (B,N,C)
        # labels - (B,N,1)
        point_embedding = self.pe_layer.forward_with_coords(points, self.input_image_size)
        point_embedding[labels == -1] = 0.0
        point_embedding[labels == -1] += self.not_a_point_embed.weight
        point_embedding[labels == 0] += self.point_embeddings[0].weight
        point_embedding[labels == 1] += self.point_embeddings[1].weight
        return point_embedding

    def _embed_boxes(
            self,
            boxes: torch.Tensor,
    ) -> torch.Tensor:
        """Embeds box prompts."""
        boxes = boxes + 0.5  # Shift to center of pixel
        coords = boxes.reshape(-1, 2, 2)
        corner_embedding = self.pe_layer.forward_with_coords(coords, self.input_image_size)
        corner_embedding[:, 0, :] += self.point_embeddings[2].weight
        corner_embedding[:, 1, :] += self.point_embeddings[3].weight
        return corner_embedding

    def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor:
        """Embeds mask prompts."""
        mask_embedding = self.mask_downscaling(masks)
        return mask_embedding


class PositionEmbeddingRandom(nn.Module):
    """
    Positional encoding using random spatial frequencies.
    """

    def __init__(
            self,
            num_pos_feats: int = 64,
            scale: Optional[float] = None
    ) -> None:
        """

        :param num_pos_feats:位置编码的特征维度的数量
        :param scale:
        """
        super().__init__()
        if scale is None or scale <= 0.0:
            scale = 1.0
        self.register_buffer(
            "positional_encoding_gaussian_matrix",
            scale * torch.randn((2, num_pos_feats)),
        )

    def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
        """
        对坐标进行位置编码
        :param coords(torch.Tensor):已被归一化，shape是(h,w,2)
        :return: shape 是 (h,w,embed_dim)
        """
        coords = 2 * coords - 1  # 这里的coords原本已经归一化到了[0,1]，计算后为[-1,1]
        # positional_encoding_gaussian_matrix - (2,embed_dim//2)
        coords = coords @ self.positional_encoding_gaussian_matrix
        coords = 2 * np.pi * coords
        return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)

    def forward(self, size: Tuple[int, int]) -> torch.Tensor:
        """Generate positional encoding for a grid of the specified size."""
        h, w = size
        device: Any = self.positional_encoding_gaussian_matrix.device
        grid = torch.ones((h, w), device=device, dtype=torch.float32)  # one就是每个元素都是1，grid — （h,w,1)
        y_embed = grid.cumsum(dim=0) - 0.5
        x_embed = grid.cumsum(dim=1) - 0.5
        y_embed = y_embed / h
        x_embed = x_embed / w
        pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))
        return pe.permute(2, 0, 1)  # Embed_dim x H x W

    def forward_with_coords(
            self, coords_input: torch.Tensor, image_size: Tuple[int, int]
    ) -> torch.Tensor:
        """
        :param coords_input(torch.Tensor): shape为(B,N,2)的张量，其中B表示批量数，N表示点的个数
        :param image_size:
        :return:
        """
        coords = coords_input.clone()
        # 归一化
        coords[:, :, 0] = coords[:, :, 0] / image_size[1]  # 对x方向上的坐标进行归一化
        coords[:, :, 1] = coords[:, :, 1] / image_size[0]  # 对y方向上的坐标进行归一化
        return self._pe_encoding(coords.to(torch.float))  # B x N x C
